import os
from docx import Document
from docx.shared import Inches
from PIL import Image
import pyheif


def add_heic_images_to_docx(folder_path, docx_filename):
    doc = Document()
    doc.add_heading('HEIC Images', level=1)

    # Iterate over files in the folder
    for filename in os.listdir(folder_path):
        if filename.endswith('.heic'):
            image_path = os.path.join(folder_path, filename)
            heif_file = pyheif.read(image_path)
            image = Image.frombytes(
                heif_file.mode,
                heif_file.size,
                heif_file.data,
                "raw",
                heif_file.mode,
                heif_file.stride,
            )
            width, height = image.size

            # Add the image to the document
            doc.add_picture(image_path, width=Inches(width / 96), height=Inches(height / 96))

    # Save the document
    doc.save(docx_filename)


# Example usage
if __name__ == "__main__":
    folder_path = "/Users/zhangxiangbin/Desktop/KID/"  # Path to the folder containing HEIC images
    docx_filename = "output.docx"  # Name for the output Word document
    add_heic_images_to_docx(folder_path, docx_filename)
